//! 控制流指令编码模块
//! 
//! 实现 8086 控制流指令：JMP, CALL, RET, 条件跳转

use crate::ast::Operand;
use crate::error::{AssemblerError, Result};
use super::registers::RegisterEncoding;

/// 控制流指令编码器
pub struct ControlFlowEncoder;

impl ControlFlowEncoder {
    /// 编码控制流指令
    pub fn encode(&self, instruction: &str, operands: &[Operand]) -> Result<Vec<u8>> {
        match instruction {
            "跳转" => self.encode_jmp(operands),
            "调用" => self.encode_call(operands),
            "返回" => self.encode_ret(operands),
            // 有符号比较跳转（已实现）
            "等于跳转" => self.encode_je(operands),
            "不等于跳转" => self.encode_jne(operands),
            "大于跳转" => self.encode_jg(operands),
            "大于等于跳转" => self.encode_jge(operands),
            "小于跳转" => self.encode_jl(operands),
            "小于等于跳转" => self.encode_jle(operands),
            // 基于标志位的跳转
            "进位跳转" => self.encode_conditional_jump(operands, 0x72, "进位跳转"),
            "无进位跳转" => self.encode_conditional_jump(operands, 0x73, "无进位跳转"),
            "溢出跳转" => self.encode_conditional_jump(operands, 0x70, "溢出跳转"),
            "无溢出跳转" => self.encode_conditional_jump(operands, 0x71, "无溢出跳转"),
            "零跳转" => self.encode_conditional_jump(operands, 0x74, "零跳转"),
            "非零跳转" => self.encode_conditional_jump(operands, 0x75, "非零跳转"),
            "符号跳转" => self.encode_conditional_jump(operands, 0x78, "符号跳转"),
            "无符号跳转" => self.encode_conditional_jump(operands, 0x79, "无符号跳转"),
            "偶跳转" => self.encode_conditional_jump(operands, 0x7A, "偶跳转"),
            "奇跳转" => self.encode_conditional_jump(operands, 0x7B, "奇跳转"),
            // 无符号比较跳转
            "高于跳转" => self.encode_conditional_jump(operands, 0x77, "高于跳转"),
            "高于等于跳转" => self.encode_conditional_jump(operands, 0x73, "高于等于跳转"),
            "低于跳转" => self.encode_conditional_jump(operands, 0x72, "低于跳转"),
            "低于等于跳转" => self.encode_conditional_jump(operands, 0x76, "低于等于跳转"),
            // LOOP 系列指令
            "循环" => self.encode_loop(operands),
            "等于循环" => self.encode_loope(operands),
            "不等于循环" => self.encode_loopne(operands),
            _ => Err(AssemblerError::CodeGenError {
                message: format!("未知的控制流指令: {}", instruction),
            }),
        }
    }

    /// 编码跳转指令
    fn encode_jmp(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 1 {
            return Err(AssemblerError::OperandError {
                message: "跳转 指令需要 1 个操作数".to_string(),
                line: 0,
            });
        }

        match &operands[0] {
            Operand::Immediate(offset) => {
                // 相对跳转：E9 = JMP rel16
                // offset 已经是相对偏移（由代码生成器计算）
                let offset_i16 = *offset as i16;
                let offset_u16 = offset_i16 as u16;
                Ok(vec![0xE9, offset_u16 as u8, (offset_u16 >> 8) as u8])
            }
            Operand::Label(_) => {
                // 标签应该在代码生成器中已经转换为相对偏移
                Err(AssemblerError::OperandError {
                    message: "内部错误：标签未转换为地址".to_string(),
                    line: 0,
                })
            }
            _ => Err(AssemblerError::OperandError {
                message: "跳转 指令只支持标签或立即数".to_string(),
                line: 0,
            }),
        }
    }

    /// 编码等于跳转指令（JE）
    fn encode_je(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        self.encode_conditional_jump(operands, 0x74, "等于跳转")
    }

    /// 编码不等于跳转指令（JNE）
    fn encode_jne(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        self.encode_conditional_jump(operands, 0x75, "不等于跳转")
    }

    /// 编码大于跳转指令（JG）
    fn encode_jg(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        self.encode_conditional_jump(operands, 0x7F, "大于跳转")
    }

    /// 编码大于等于跳转指令（JGE）
    fn encode_jge(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        self.encode_conditional_jump(operands, 0x7D, "大于等于跳转")
    }

    /// 编码小于跳转指令（JL）
    fn encode_jl(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        self.encode_conditional_jump(operands, 0x7C, "小于跳转")
    }

    /// 编码小于等于跳转指令（JLE）
    fn encode_jle(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        self.encode_conditional_jump(operands, 0x7E, "小于等于跳转")
    }

    /// 条件跳转指令的通用编码函数
    fn encode_conditional_jump(
        &self,
        operands: &[Operand],
        opcode: u8,
        instruction_name: &str,
    ) -> Result<Vec<u8>> {
        if operands.len() != 1 {
            return Err(AssemblerError::OperandError {
                message: format!("{} 指令需要 1 个操作数", instruction_name),
                line: 0,
            });
        }

        match &operands[0] {
            Operand::Immediate(offset) => {
                // 条件跳转使用 8 位相对偏移（-128 到 +127）
                let offset_i8 = *offset as i8;
                let offset_u8 = offset_i8 as u8;
                
                // 检查偏移范围
                if *offset < -128 || *offset > 127 {
                    return Err(AssemblerError::OperandError {
                        message: format!(
                            "{} 偏移超出范围: {} 字节（条件跳转只支持 -128 到 +127 字节）",
                            instruction_name, offset
                        ),
                        line: 0,
                    });
                }
                
                Ok(vec![opcode, offset_u8])
            }
            Operand::Label(_) => {
                // 标签应该在代码生成器中已经转换为相对偏移
                Err(AssemblerError::OperandError {
                    message: "内部错误：标签未转换为地址".to_string(),
                    line: 0,
                })
            }
            _ => Err(AssemblerError::OperandError {
                message: format!("{} 指令只支持标签或立即数", instruction_name),
                line: 0,
            }),
        }
    }

    /// 编码调用指令（CALL）
    fn encode_call(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.len() != 1 {
            return Err(AssemblerError::OperandError {
                message: "调用 指令需要 1 个操作数".to_string(),
                line: 0,
            });
        }

        match &operands[0] {
            Operand::Immediate(offset) => {
                // CALL rel16：E8 + 16位相对偏移
                // offset 已经是相对偏移（由代码生成器计算）
                let offset_i16 = *offset as i16;
                let offset_u16 = offset_i16 as u16;
                Ok(vec![0xE8, offset_u16 as u8, (offset_u16 >> 8) as u8])
            }
            Operand::Register(reg) => {
                // CALL reg：FF /2 (reg字段=010表示CALL)
                let reg_code = RegisterEncoding::get_16bit_register_code(reg)
                    .ok_or_else(|| AssemblerError::OperandError {
                        message: format!("无效的16位寄存器: {}", reg),
                        line: 0,
                    })?;
                let opcode = 0xFF;
                let modrm = (0b11 << 6) | (0b010 << 3) | reg_code; // reg字段=010表示CALL
                Ok(vec![opcode, modrm])
            }
            Operand::Label(_) => {
                // 标签应该在代码生成器中已经转换为相对偏移
                Err(AssemblerError::OperandError {
                    message: "内部错误：标签未转换为地址".to_string(),
                    line: 0,
                })
            }
            _ => Err(AssemblerError::OperandError {
                message: "调用 指令只支持标签、立即数或寄存器".to_string(),
                line: 0,
            }),
        }
    }

    /// 编码返回指令（RET）
    fn encode_ret(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        if operands.is_empty() {
            // RET：C3（1字节）
            Ok(vec![0xC3])
        } else if operands.len() == 1 {
            // RET imm16：C2 + 16位立即数（3字节）
            if let Operand::Immediate(imm) = &operands[0] {
                if *imm < 0 || *imm > 65535 {
                    return Err(AssemblerError::OperandError {
                        message: format!("返回 立即数必须在 0 到 65535 之间，得到: {}", imm),
                        line: 0,
                    });
                }
                let imm_u16 = *imm as u16;
                Ok(vec![0xC2, imm_u16 as u8, (imm_u16 >> 8) as u8])
            } else {
                Err(AssemblerError::OperandError {
                    message: "返回 指令的立即数操作数必须是数字".to_string(),
                    line: 0,
                })
            }
        } else {
            Err(AssemblerError::OperandError {
                message: "返回 指令最多需要 1 个操作数".to_string(),
                line: 0,
            })
        }
    }

    /// 编码循环指令（LOOP）
    /// 
    /// LOOP label: CX--, if CX != 0 then jump to label
    /// 操作码: 0xE2
    fn encode_loop(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        self.encode_conditional_jump(operands, 0xE2, "循环")
    }

    /// 编码等于循环指令（LOOPE/LOOPZ）
    /// 
    /// LOOPE label: CX--, if CX != 0 && ZF=1 then jump to label
    /// 操作码: 0xE1
    fn encode_loope(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        self.encode_conditional_jump(operands, 0xE1, "等于循环")
    }

    /// 编码不等于循环指令（LOOPNE/LOOPNZ）
    /// 
    /// LOOPNE label: CX--, if CX != 0 && ZF=0 then jump to label
    /// 操作码: 0xE0
    fn encode_loopne(&self, operands: &[Operand]) -> Result<Vec<u8>> {
        self.encode_conditional_jump(operands, 0xE0, "不等于循环")
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_encode_jmp_immediate() {
        let encoder = ControlFlowEncoder;
        let operands = vec![Operand::Immediate(0x0100)];
        let result = encoder.encode("跳转", &operands).unwrap();
        // JMP rel16 = E9 00 01
        assert_eq!(result, vec![0xE9, 0x00, 0x01]);
    }

    #[test]
    fn test_encode_je() {
        let encoder = ControlFlowEncoder;
        let operands = vec![Operand::Immediate(10)];
        let result = encoder.encode("等于跳转", &operands).unwrap();
        // JE rel8 = 74 0A
        assert_eq!(result, vec![0x74, 0x0A]);
    }

    #[test]
    fn test_encode_jne() {
        let encoder = ControlFlowEncoder;
        let operands = vec![Operand::Immediate(-10)];
        let result = encoder.encode("不等于跳转", &operands).unwrap();
        // JNE rel8 = 75 F6
        assert_eq!(result, vec![0x75, 0xF6]);
    }

    #[test]
    fn test_encode_call_register() {
        let encoder = ControlFlowEncoder;
        let operands = vec![Operand::Register("基址".to_string())];
        let result = encoder.encode("调用", &operands).unwrap();
        // CALL BX = FF D3
        assert_eq!(result, vec![0xFF, 0xD3]);
    }

    #[test]
    fn test_encode_ret() {
        let encoder = ControlFlowEncoder;
        let operands = vec![];
        let result = encoder.encode("返回", &operands).unwrap();
        // RET = C3
        assert_eq!(result, vec![0xC3]);
    }

    #[test]
    fn test_encode_jmp_wrong_operand() {
        let encoder = ControlFlowEncoder;
        let operands = vec![Operand::Register("累加器".to_string())];
        let result = encoder.encode("跳转", &operands);
        assert!(result.is_err());
    }

    #[test]
    fn test_encode_jc() {
        let encoder = ControlFlowEncoder;
        let operands = vec![Operand::Immediate(10)];
        let result = encoder.encode("进位跳转", &operands).unwrap();
        // JC rel8 = 72 0A
        assert_eq!(result, vec![0x72, 0x0A]);
    }

    #[test]
    fn test_encode_jnc() {
        let encoder = ControlFlowEncoder;
        let operands = vec![Operand::Immediate(-10)];
        let result = encoder.encode("无进位跳转", &operands).unwrap();
        // JNC rel8 = 73 F6
        assert_eq!(result, vec![0x73, 0xF6]);
    }

    #[test]
    fn test_encode_jo() {
        let encoder = ControlFlowEncoder;
        let operands = vec![Operand::Immediate(20)];
        let result = encoder.encode("溢出跳转", &operands).unwrap();
        // JO rel8 = 70 14
        assert_eq!(result, vec![0x70, 0x14]);
    }

    #[test]
    fn test_encode_ja() {
        let encoder = ControlFlowEncoder;
        let operands = vec![Operand::Immediate(30)];
        let result = encoder.encode("高于跳转", &operands).unwrap();
        // JA rel8 = 77 1E
        assert_eq!(result, vec![0x77, 0x1E]);
    }

    #[test]
    fn test_encode_jb() {
        let encoder = ControlFlowEncoder;
        let operands = vec![Operand::Immediate(40)];
        let result = encoder.encode("低于跳转", &operands).unwrap();
        // JB rel8 = 72 28
        assert_eq!(result, vec![0x72, 0x28]);
    }

    #[test]
    fn test_encode_loop() {
        let encoder = ControlFlowEncoder;
        let operands = vec![Operand::Immediate(10)];
        let result = encoder.encode("循环", &operands).unwrap();
        // LOOP rel8 = E2 0A
        assert_eq!(result, vec![0xE2, 0x0A]);
    }

    #[test]
    fn test_encode_loope() {
        let encoder = ControlFlowEncoder;
        let operands = vec![Operand::Immediate(-10)];
        let result = encoder.encode("等于循环", &operands).unwrap();
        // LOOPE rel8 = E1 F6
        assert_eq!(result, vec![0xE1, 0xF6]);
    }

    #[test]
    fn test_encode_loopne() {
        let encoder = ControlFlowEncoder;
        let operands = vec![Operand::Immediate(20)];
        let result = encoder.encode("不等于循环", &operands).unwrap();
        // LOOPNE rel8 = E0 14
        assert_eq!(result, vec![0xE0, 0x14]);
    }
}

